#!/usr/bin/env python3
# coding=utf-8

import rospy
from std_msgs.msg import String
from openai import OpenAI

def cbQuestion(msg):
    rospy.loginfo("--------------------")
    rospy.loginfo("收到问题: %s", msg.data)

    global conversation, msg_log
    if conversation:
        msg_log.append({"role": "user", "content": msg.data})
    else:
        msg_log = [{"role": "user", "content": msg.data}]

    # 每次请求时创建新的客户端实例
    client = OpenAI(api_key=api_key, base_url="https://api.deepseek.com")
    try:
        # 使用新版 ChatCompletion 接口调用 API
        response = client.chat.completions.create(
            model=model_engine,
            messages=msg_log,
            temperature=0.7,
            max_tokens=1024,
            stream=False
        )

        # 解析返回结果
        response_text = response.choices[0].message.content.strip()
        rospy.logwarn("回复: %s", response_text)

        if conversation:
            msg_log.append({"role": "assistant", "content": response_text})

        # 发布返回结果
        answer_msg = String()
        answer_msg.data = response_text
        response_pub.publish(answer_msg)

    except Exception as e:
        rospy.logerr("与 DeepSeek 通信出错: %s", str(e))

if __name__ == "__main__":
    rospy.init_node("chatgpt_node")

    # 从 ROS 参数服务器读取配置
    api_key = rospy.get_param('~deepseek/api_key')
    model_engine = rospy.get_param('~deepseek/model', "deepseek-chat")
    personality = rospy.get_param('~deepseek/personality', "")
    conversation = rospy.get_param('~deepseek/conversation', False)

    rospy.logwarn("DeepSeek: 当前使用模型 %s", model_engine)

    question_sub = rospy.Subscriber("/wpr_ask", String, cbQuestion, queue_size=1)
    response_pub = rospy.Publisher("/chatspt_answer", String, queue_size=1)

    msg_log = []

    # 如果配置了 personality，则在消息记录中添加系统信息
    if personality:
        msg_log.append({"role": "system", "content": f"你是一名{personality}"})
        rospy.logwarn(f"DeepSeek: 我是一名{personality}，已经准备好回答您的问题！")
    else:
        rospy.logwarn("DeepSeek: 已经准备好回答您的问题！")

    rospy.spin()
